#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""计算栅格点在平面上的欧氏距离

对于平面上给定的栅格点, 计算平面上其他点 (NoData 的点) 距离这些点的距离。

Example:
    $ python euclidean-distance-raster.py points.tif out.tif
"""

import sys
import math
import numpy as np
from osgeo import gdal, osr

def compute_distance(point_a, point_b):
    a_x, a_y = point_a
    b_x, b_y = point_b
    distance = math.sqrt(math.pow(b_x - a_x, 2) + math.pow(b_y - a_y, 2))

    return distance

def find_minial_distance(target_point, coordinates):
    p_x, p_y = target_point
    
    minimal = None
    for point in coordinates:
        distance = compute_distance(target_point, point)

        if minimal is None:
            minimal = distance
        elif distance < minimal:
            minimal = distance

    return minimal

def raster2array(raster_fn):
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)
    no_data = band.GetNoDataValue()
    array = band.ReadAsArray()

    return array, no_data

def array2raster(new_raster_fn, raster_fn, array):
    raster = gdal.Open(raster_fn)
    geotransform = raster.GetGeoTransform()
    origin_x = geotransform[0]
    origin_y = geotransform[3]
    pixel_width = geotransform[1]
    pixel_height = geotransform[5]
    rows, cols = array.shape

    driver = gdal.GetDriverByName('GTiff')
    out_raster = driver.Create(new_raster_fn, cols, rows, 1, gdal.GDT_Float32)
    out_raster.SetGeoTransform((origin_x, pixel_width, 0, origin_y, 0, pixel_height))
    outband = out_raster.GetRasterBand(1)
    outband.WriteArray(array)
    out_raster_srs = osr.SpatialReference()
    out_raster_srs.ImportFromWkt(raster.GetProjectionRef())
    out_raster.SetProjection(out_raster_srs.ExportToWkt())
    outband.FlushCache()

def data_point(array, no_data):
    coordinates = []
    for item in array[array != no_data]:
        index = np.where(array == item)
        coordinates.append(list(zip(index[0], index[1])))

    return coordinates

def pixel2geo(raster_fn, point):
    raster = gdal.Open(raster_fn)
    band = raster.GetRasterBand(1)

    geotransform = raster.GetGeoTransform()
    originX = geotransform[0]
    originY = geotransform[3]
    pixelWidth = geotransform[1]
    pixelHeight = geotransform[5]

    x_offset, y_offset = point

    x = x_offset * pixelWidth + originX
    y = y_offset * pixelHeight + originY

    return (x, y)

def main(raster_fn, new_raster_fn):
    array, no_data = raster2array('source_PointToRaster.tif')
    index = data_point(array, no_data)

    distance_array = np.zeros_like(array)

    coordinates = []
    for items in index:
        for item in items:
            point = pixel2geo(raster_fn, item)
            coordinates.append(point)

    for x, y in np.ndindex(array.shape):
        geo_point = pixel2geo(raster_fn, (x, y))
        distance_array[x, y] = find_minial_distance(geo_point, coordinates)
    
    array2raster(new_raster_fn, raster_fn, distance_array)

if __name__ == "__main__":
    raster_fn = None
    new_raster_fn = None

    argv = sys.argv

    i = 1
    while i < len(argv):
        arg = argv[i]

        if raster_fn is None:
            raster_fn = argv[i]

        elif new_raster_fn is None:
            new_raster_fn = argv[i]

        else:
            print("python script.py point.tif out.tif")
            sys.exit(1)

        i = i + 1

    if new_raster_fn is None:
        new_raster_fn = 'OUT.tif'

    main(raster_fn, new_raster_fn)

